# 输入model路径，导出engine格式的模型到输出目录
from ultralytics import YOLO
import os
import cv2
import numpy as np

os.environ["CUDA_MODULE_LOADING"] = "LAZY"
model_path = "/home/xlli67/Project/model_export/input/best.pt"
output_dir = "/home/xlli67/Project/model_export/output"

os.makedirs(output_dir, exist_ok=True)

try:
    model = YOLO(model_path)
    model.export(format="engine", 
                dynamic=False, 
                batch=1,
                imgsz=640,
                workspace=8,
                half=True)
    print("Exported to TensorRT engine format successfully!")
except Exception as e:
    print(f"Error during export: {e}")
